package com.util;

import com.example.Category;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.*;
import java.util.function.*;
import java.util.stream.Collectors;

/**
 * tree和list互转工具类
 */
public class TreeUtils {

    /**
     * list转tree
     */
    public static <T, R> List<T> list2Tree(List<T> source, Function<T, R> idFn, Function<T, R> pidFn,
                                           Predicate<T> isRoot, BiConsumer<T, List<T>> setChildren) {
        return list2Tree(source, idFn, pidFn, isRoot, setChildren, null);
    }

    /**
     * list转tree
     * 代码最简单的版本，使用了两次循环
     *
     * @param source      原始列表
     * @param idFn        获取id的方法
     * @param pidFn       获取父id的方法
     * @param isRoot      判断是否是根节点的方法
     * @param setChildren 设置子节点的方法
     * @param setIdx      设置节点层级的方法（可选的操作，可为null）
     * @param <T>         元素类型
     * @param <R>         id类型
     * @return tree
     */
    public static <T, R> List<T> list2Tree(List<T> source, Function<T, R> idFn, Function<T, R> pidFn,
                                           Predicate<T> isRoot, BiConsumer<T, List<T>> setChildren,
                                           BiConsumer<T, Integer> setIdx) {
        // 1.找出根节点，并将子节点根据父id分组
        List<T> rootList = new ArrayList<>();
        Map<R, List<T>> nodeMap = new HashMap<>();
        for (T t : source) {
            if (isRoot.test(t)) {
                rootList.add(t);
            } else {
                List<T> nodeList = nodeMap.getOrDefault(pidFn.apply(t), new ArrayList<>());
                nodeList.add(t);
                nodeMap.put(pidFn.apply(t), nodeList);
            }
        }
        // 遍历根节点，递归设置子节点
        rootList.forEach(e -> buildTree(e, nodeMap, idFn, setChildren, setIdx, 0));
        return rootList;
    }

    /**
     * 递归构造tree的子节点
     *
     * @param current     当前节点
     * @param nodeMap     节点集合
     * @param idFn        获取id的方法
     * @param setChildren 设置子节点的方法
     * @param setIdx      设置节点层级的方法（可选的操作，可为null）
     * @param idx         当前节点的层级（start from 0）
     * @param <T>         元素类型
     * @param <R>         id类型
     */
    private static <T, R> void buildTree(T current, Map<R, List<T>> nodeMap,
                                         Function<T, R> idFn, BiConsumer<T, List<T>> setChildren,
                                         BiConsumer<T, Integer> setIdx, int idx) {
        // 获取当前元素的子节点
        List<T> children = nodeMap.get(idFn.apply(current));
        // 设置子节点
        setChildren.accept(current, children);
        // 递归
        if (children != null && children.size() > 0) {
            children.forEach(e -> buildTree(e, nodeMap, idFn, setChildren, setIdx, idx + 1));
        }
        // 设置当前节点的层级
        Optional.ofNullable(setIdx).ifPresent(l -> l.accept(current, idx));
    }

    /**
     * tree转list
     *
     * @param source        源数据
     * @param getChildrenFn 获取子节点的方法
     * @param <T>           元素类型
     * @return list
     */
    public static <T> List<T> tree2List(List<T> source, Function<T, List<T>> getChildrenFn) {
        List<T> target = new ArrayList<>();
        loopTree(source, getChildrenFn, target::add, null);
        return target;
    }

    /**
     * 递归遍历tree
     * 以深度优先的方式
     *
     * @param source        源数据
     * @param getChildrenFn 获取子节点的方法
     * @param preListener   前置监听器，用于在遍历节点前对其做自定义的操作，可为null
     * @param postListener  后置监听器，用于在遍历节点后对其做自定义的操作，可为null
     * @param <T>           元素类型
     */
    public static <T> void loopTree(List<T> source, Function<T, List<T>> getChildrenFn,
                                    Consumer<T> preListener, Consumer<T> postListener) {
        // 递归中止条件
        if (source == null || source.size() == 0) {
            return;
        }
        // 对当前层级的节点进行递归查找
        source.forEach(e -> {
            Optional.ofNullable(preListener).ifPresent(l -> l.accept(e));
            // 获取子节点
            List<T> children = getChildrenFn.apply(e);
            // 递归
            loopTree(children, getChildrenFn, preListener, postListener);
            Optional.ofNullable(postListener).ifPresent(l -> l.accept(e));
        });
    }

    public static void main(String[] args) throws JsonProcessingException {
        ArrayList<Category> categories = new ArrayList<>();
        categories.add(Category.builder().id(1L).pid(null).name("数码产品").build());
        categories.add(Category.builder().id(2L).pid(null).name("洗护用品").build());
        categories.add(Category.builder().id(3L).pid(1L).name("手机").build());
        categories.add(Category.builder().id(4L).pid(1L).name("电脑").build());
        categories.add(Category.builder().id(5L).pid(2L).name("纸巾").build());
        categories.add(Category.builder().id(6L).pid(2L).name("毛巾").build());
        categories.add(Category.builder().id(7L).pid(3L).name("安卓手机").build());
        categories.add(Category.builder().id(8L).pid(3L).name("苹果手机").build());

        ObjectMapper objectMapper = new ObjectMapper();

        List<Category> tree = TreeUtils.list2Tree(categories, Category::getId, Category::getPid,
                e -> e.getPid() == null, Category::setChildren, Category::setIndex);
        System.out.println(objectMapper.writeValueAsString(tree));

        List<Category> list = TreeUtils.tree2List(tree, Category::getChildren);
        System.out.println(objectMapper.writeValueAsString(list));
    }
}
